package org.oregan.gui.capi;

import org.oregan.asn1.ASN1;
import org.oregan.asn1.DEREncoder;
import org.oregan.asn1.EncodableTLV;
import org.oregan.asn1.X509CertificateHolder;

import javax.swing.JProgressBar;
import javax.swing.SwingUtilities;
import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.security.DigestInputStream;
import java.security.InvalidKeyException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.Signature;
import java.security.SignatureException;
import java.security.UnrecoverableKeyException;
import java.security.cert.CertificateEncodingException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.Date;
import java.util.jar.JarEntry;
import java.util.jar.JarOutputStream;
import java.util.zip.CRC32;

public class XPISigner
{
    public static final String DIGEST_MANIFST_MESSAGE =
        "Name: {0}\n" +
            "Digest-Algorithms: MD5 SHA1\n" +
            "MD5-Digest: {1}\n" +
            "SHA1-Digest: {2}\n";

    public static final String DIGEST_MESSAGE =
        "Digest-Algorithms: MD5 SHA1\n" +
            "MD5-Digest: {0}\n" +
            "SHA1-Digest: {1}\n";

    static final String PERSONALISATION = "Created-by: MultiSigner\n" +
        "Multi-Signer-Version: 2.0 http://o-regan.org\n";

    static final String MANIFEST_PREAMBLE =
        "Manifest-Version: 1.0\n" +
            PERSONALISATION +
            "Comments: PLEASE DO NOT EDIT THIS FILE. YOU WILL BREAK IT.\n";

    static final String ZIGBERT_PREAMBLE =
        "Signature-Version: 1.0\n" +
            PERSONALISATION +
            "Comments: PLEASE DO NOT EDIT THIS FILE. YOU WILL BREAK IT.\n";


    public static final String META_INF_ZIGBERT_RSA = "META-INF/zigbert.rsa";
    public static final String META_INF_ZIGBERT_SF = "META-INF/zigbert.sf";
    public static final String META_INF_MANIFEST_MF = "META-INF/manifest.mf";


    private MessageDigest digestSHA = null;
    private MessageDigest digestMD5 = null;

    private StringBuffer manifest;
    private StringBuffer zigbert;
    private File outputFile;
    private String pfxfile;
    private String password;

    private ArrayList listing;

    private File baseDir;
    private String alias;
    private ArrayList<FileInfo> jarOrder;

    public XPISigner(String alias, File baseDir, ArrayList<String> listing, File outputFile)
    {
        this.alias = alias;
        this.baseDir = baseDir;
        this.listing = listing;
        this.outputFile = outputFile;
        try
        {
            digestSHA = MessageDigest.getInstance("SHA1");
            digestMD5 = MessageDigest.getInstance("MD5");
        } catch (NoSuchAlgorithmException e)
        {
            //
        }


    }

     public XPIInfo archiveXPI(JProgressBar prog) throws XPIException
    {
        try
        {

            jarOrder.add(0, inspectFile(META_INF_ZIGBERT_RSA, false)); // rsa file at start and uncompressed.
            jarOrder.add(inspectFile(META_INF_MANIFEST_MF));
            jarOrder.add(inspectFile(META_INF_ZIGBERT_SF));


            XPIInfo info = new XPIInfo(outputFile ,"", jarOrder.size());

            JarOutputStream jos = new JarOutputStream(new FileOutputStream(outputFile));
            updateMaxValue(prog, jarOrder.size());
            int i=1;
            for (FileInfo fi : jarOrder)
            {
                jarFile(jos, fi);
                updateProgress(prog,i++);
            }
            updateProgress(prog, jarOrder.size());
            jos.close();

            return info;

        } catch (IOException e)
        {
            //
        }
        return null;
    }

    public void compileManifests(JProgressBar prog)
        throws IOException, XPIException
    {
        manifest = new StringBuffer();
        zigbert = new StringBuffer();

        jarOrder = new ArrayList<FileInfo>();

        updateMaxValue(prog,listing.size());
        updateProgress(prog,1);


        startManifestEntries();
        int i=1;
        for (Object o : listing)
        {
            String line = (String) o;
            try
            {
                FileInfo metadata = inspectFile(line);
                jarOrder.add(metadata);
                addManifestEntry(metadata);
                updateProgress(prog, i++);
            } catch (XPIException e)
            {
                System.err.println(e.getMessage());
            }
        }

        File f = new File(baseDir, META_INF_MANIFEST_MF);
        if (!f.getParentFile().exists())
            f.getParentFile().mkdirs();

        Utils.saveMessage((manifest.toString().trim() + "\n").getBytes("LATIN1"), f);

        byte[] zbytes = (zigbert.toString().trim() + "\n").getBytes("LATIN1");
        f = new File(baseDir, META_INF_ZIGBERT_SF);
        Utils.saveMessage(zbytes, f);


        byte[] signature = sign(zbytes, alias);

        f = new File(baseDir, META_INF_ZIGBERT_RSA);
        Utils.saveMessage(signature, f);
        updateProgress(prog,listing.size());
    }




    FileInfo inspectFile(String filePath) throws XPIException
    {
        return inspectFile(filePath, true);
    }

    /**
     * Inspect the file and build a fileinfo object
     *
     * @param filePath
     * @param compress
     * @return
     */
    private FileInfo inspectFile(String filePath, boolean compress) throws XPIException
    {
        long length = 0;
        long crc = 0;
        byte[] md5 = new byte[0];
        byte[] sha1 = new byte[0];

        File tmp = new File(baseDir, filePath);
        if (!tmp.exists())
            throw new XPIException("  File \'" + tmp.getAbsolutePath() + "\' does not exist.", Constants.ERR_FILE_NOT_FOUND);

        try
        {

            length = tmp.length();

            BufferedInputStream bis = null;
            try
            {
                bis = new BufferedInputStream(new FileInputStream(tmp));
            } catch (FileNotFoundException e1)
            {
                // handled above
            }

            digestMD5.reset();
            digestSHA.reset();

            DigestInputStream md5Stream = new DigestInputStream(bis, digestMD5);
            DigestInputStream shaStream = new DigestInputStream(md5Stream, digestSHA);

            CRC32 crc32 = new CRC32();

            byte[] data = new byte[1024 * 2];
            int byteCount;
            while ((byteCount = shaStream.read(data)) > -1)
            {
                crc32.update(data, 0, byteCount);
            }
            crc = crc32.getValue();
            crc32.reset();
            md5 = md5Stream.getMessageDigest().digest();
            sha1 = shaStream.getMessageDigest().digest();
            bis.close();
        } catch (IOException e1)
        {
            throw new XPIException("Error reading from \'" + tmp + "\'", Constants.ERR_ERROR_READING_FILE);
        }


        return new FileInfo(filePath, length, crc, md5, sha1, compress);

    }

    void startManifestEntries()
    {
        digestMD5.reset();
        digestSHA.reset();

        byte[] bytes = null;
        try
        {
            bytes = MANIFEST_PREAMBLE.getBytes("LATIN1");
        } catch (UnsupportedEncodingException e)
        {
            //
        }
        byte[] md5 = digestMD5.digest(bytes);
        byte[] sha = digestSHA.digest(bytes);

        String md5Text = Utils.toB64(md5);
        String shaText = Utils.toB64(sha);

        String zigbertR = MessageFormat.format(DIGEST_MESSAGE, md5Text, shaText);

        manifest.append(MANIFEST_PREAMBLE).append("\n");
        zigbert.append(ZIGBERT_PREAMBLE).append(zigbertR).append("\n");
    }

    void addManifestEntry(FileInfo fileInfo)
    {

        String name = fileInfo.getName();


        String md5Text = Utils.toB64(fileInfo.getMd5());
        String sha1Text = Utils.toB64(fileInfo.getSha1());


        String mfest = MessageFormat.format(DIGEST_MANIFST_MESSAGE, name, md5Text, sha1Text);

        manifest.append(mfest).append("\n");

        digestMD5.reset();
        digestSHA.reset();

        try
        {
            byte[] bytes = mfest.getBytes("LATIN1");
            md5Text = Utils.toB64(digestMD5.digest(bytes));
            sha1Text = Utils.toB64(digestSHA.digest(bytes));

            String ziggy = MessageFormat.format(DIGEST_MANIFST_MESSAGE, name, md5Text, sha1Text);
            zigbert.append(ziggy).append("\n");

        } catch (UnsupportedEncodingException e)
        {
            e.printStackTrace();
        }
    }

    void jarFile(JarOutputStream stream, FileInfo metadata)
    {
        try
        {
            String filePath = metadata.getName();

            File f = new File(baseDir, filePath);

            BufferedInputStream bis = new BufferedInputStream(new FileInputStream(f));

            JarEntry fileEntry = new JarEntry(filePath);

            if (metadata.isCompressed())
            {
                fileEntry.setMethod(JarEntry.DEFLATED);
            } else
            {
                fileEntry.setMethod(JarEntry.STORED);
                fileEntry.setSize(metadata.getLength());
                fileEntry.setCompressedSize(metadata.getLength());
                fileEntry.setCrc(metadata.getCRC());
            }
            stream.putNextEntry(fileEntry);
            Utils.slurp(bis, stream);
            stream.closeEntry();
            bis.close();

        } catch (IOException e)
        {
            e.printStackTrace();
        }
    }


    public byte[] sign(byte[] zbytes, String alias) throws XPIException
    {
        try
        {
            KeyStore msCertStore = KeyStore.getInstance("Windows-MY", "SunMSCAPI");
            msCertStore.load(null, null);

            PrivateKey capiKeyRef = (PrivateKey) msCertStore.getKey(alias, null);
            X509Certificate[] chain = (X509Certificate[]) msCertStore.getCertificateChain(alias);


            EncodableTLV seq = new EncodableTLV(ASN1.SEQUENCE);
            seq.addChild(new EncodableTLV(ASN1.OBJECTIDENTIFIER, ASN1.OIDsSignedData));
            EncodableTLV ctx = new EncodableTLV((byte) 0xA0); //[0]
            seq.addChild(ctx);
            EncodableTLV seq2 = new EncodableTLV(ASN1.SEQUENCE);
            ctx.addChild(seq2);
            byte[] ONE = {1};
            seq2.addChild(new EncodableTLV(ASN1.INTEGER, ONE)); // Version
            // AlgID
            EncodableTLV algID = algorithmIdentifier();
            seq2.addChild(algID);
            seq2.addChild(new EncodableTLV(ASN1.SEQUENCE)).addChild(new EncodableTLV(ASN1.OBJECTIDENTIFIER, ASN1.OIDsData));

            EncodableTLV certChain = new EncodableTLV((byte) 0xA0);


            X509CertificateHolder hldr = null;
            for (int i = 0; i < chain.length; i++)
            {
                byte[] encoded = chain[i].getEncoded();
                hldr = new X509CertificateHolder("", encoded);
                EncodableTLV cert = hldr.getEncodable();
                certChain.addChild(cert);
            }
            seq2.addChild(certChain);


            EncodableTLV signerInfos = (new EncodableTLV(ASN1.SETOF));
            EncodableTLV signerInfo = new EncodableTLV(ASN1.SEQUENCE);
            signerInfo.addChild(new EncodableTLV(ASN1.INTEGER, ONE)); // Version
            // IASN

            EncodableTLV iasnSeq = new EncodableTLV(ASN1.SEQUENCE);
            hldr = new X509CertificateHolder(alias, chain[0].getEncoded());
            iasnSeq.addChild(hldr.getIssuerDN());
            iasnSeq.addChild(hldr.getSerial());

            signerInfo.addChild(iasnSeq);
            MessageDigest sha1 = MessageDigest.getInstance("SHA1");
            byte[] messageDigest = sha1.digest(zbytes);

            // HashAlg
            byte[] oid = ASN1.OIDsSHA1;
            EncodableTLV param = new EncodableTLV(ASN1.NULL);

            EncodableTLV objId = new EncodableTLV(ASN1.SEQUENCE);
            objId.addChild(new EncodableTLV(ASN1.OBJECTIDENTIFIER, oid));
            objId.addChild(param);
            signerInfo.addChild(objId);

            // AuthAttrs
            EncodableTLV authAttrs = new EncodableTLV((byte) 0xA0);

            authAttrs.addChild(attrValue(ASN1.OIDsContentType, new EncodableTLV(ASN1.OBJECTIDENTIFIER, ASN1.OIDsData)));

            authAttrs.addChild(attrValue(ASN1.OIDsSigningTime, new EncodableTLV(ASN1.UTCTIME, Utils.toUTCTime(new Date()).getBytes())));

            authAttrs.addChild(attrValue(ASN1.OIDsMessageDigest, new EncodableTLV(ASN1.OCTETSTRING, messageDigest)));

            signerInfo.addChild(authAttrs);

            // When encoding for the signature this needs to be a SETOF not a Tagged Conntext Specific
            EncodableTLV authAttrsSeq = new EncodableTLV(ASN1.SETOF);

            authAttrsSeq.addChild(attrValue(ASN1.OIDsContentType, new EncodableTLV(ASN1.OBJECTIDENTIFIER, ASN1.OIDsData)));

            authAttrsSeq.addChild(attrValue(ASN1.OIDsSigningTime, new EncodableTLV(ASN1.UTCTIME, Utils.toUTCTime(new Date()).getBytes())));

            authAttrsSeq.addChild(attrValue(ASN1.OIDsMessageDigest, new EncodableTLV(ASN1.OCTETSTRING, messageDigest)));


            byte[] authAttrDER = DEREncoder.encode(authAttrsSeq);

            sha1.reset();

            byte[] encDigest = sha1.digest(authAttrDER);

            Signature signer = Signature.getInstance("SHA1withRSA", "SunMSCAPI");
            signer.initSign(capiKeyRef);

            signer.update(authAttrDER);
            byte[] signatureBytes = signer.sign();

            // SigAlgID
            EncodableTLV algIdRSA = new EncodableTLV(ASN1.SEQUENCE);
            algIdRSA.addChild(new EncodableTLV(ASN1.OBJECTIDENTIFIER, ASN1.OIDsRSAEncryption));
            algIdRSA.addChild(new EncodableTLV(ASN1.NULL));
            // Signature bytes

            signerInfo.addChild(algIdRSA);

            signerInfo.addChild(new EncodableTLV(ASN1.OCTETSTRING, signatureBytes));

            signerInfos.addChild(signerInfo);

            seq2.addChild(signerInfos);

            byte[] out = DEREncoder.encode(seq);
            return out;
        } catch (IOException e)
        {
            e.printStackTrace();
        } catch (NoSuchProviderException e)
        {
            e.printStackTrace();
        } catch (CertificateEncodingException e)
        {
            e.printStackTrace();
        } catch (NoSuchAlgorithmException e)
        {
            e.printStackTrace();
        } catch (SignatureException e)
        {
            e.printStackTrace();
        } catch (CertificateException e)
        {
            e.printStackTrace();
        } catch (InvalidKeyException e)
        {
            e.printStackTrace();
        } catch (UnrecoverableKeyException e)
        {
            e.printStackTrace();
        } catch (KeyStoreException e)
        {
            e.printStackTrace();
        }

        return new byte[0];
    }

    private EncodableTLV algorithmIdentifier()
    {

        byte[] oid = ASN1.OIDsSHA1;
        EncodableTLV param = new EncodableTLV(ASN1.NULL);

        EncodableTLV algID = new EncodableTLV(ASN1.SETOF);
        EncodableTLV objId = new EncodableTLV(ASN1.SEQUENCE);
        objId.addChild(new EncodableTLV(ASN1.OBJECTIDENTIFIER, oid));
        objId.addChild(param);
        algID.addChild(objId);
        return algID;
    }

    private EncodableTLV attrValue(byte[] oid, EncodableTLV value)
    {
        EncodableTLV seq = new EncodableTLV(ASN1.SEQUENCE);
        seq.addChild(new EncodableTLV(ASN1.OBJECTIDENTIFIER, oid));
        seq.addChild(new EncodableTLV(ASN1.SETOF)).addChild(value);
        return seq;
    }

   
    private void updateMaxValue(final JProgressBar prog,final int max)
    {
          SwingUtilities.invokeLater(new Runnable(){

            public void run()
            {
                prog.setMaximum(max);
            }
        });
    }

     private void updateProgress(final JProgressBar prog, final int i)
    {
      SwingUtilities.invokeLater(new Runnable(){

            public void run()
            {
                prog.setValue(i);
            }
        });
    }
}
